import cv2
import numpy as np
import os
# import sys
# import time
import tkinter as tk
# from tkinter import filedialog#文件控件
from PIL import Image#图像控件
# from PIL import ImageTk#图像控件

window = tk.Tk()
window.title('人脸识别系统')
sw = window.winfo_screenwidth()#获取屏幕宽
sh = window.winfo_screenheight()#获取屏幕高
wx = 280
wh = 300
window.geometry("%dx%d+%d+%d" %(wx,wh,(sw-wx)/2 + 300,(sh-wh)/2))#窗口至指定位置
# canvas = tk.Canvas(window,bg='#c4c2c2',height=wh,width=wx) #绘制画布
# canvas.pack()

face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')


#根据给定的（x，y）坐标和宽度高度在图像上绘制矩形
def draw_rectangle(img, rect):
    (x, y, w, h) = rect
    cv2.rectangle(img, (x, y), (x + w, y + h), (128, 128, 0), 2)
# 根据给定的（x，y）坐标标识出人名


def draw_text(img, text, x, y, color):
    if(color == "GREEN"):
        cv2.putText(img, text, (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (128, 128, 0), 2)
    elif(color == "RED"):
        cv2.putText(img, text, (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)


def detect_face(img):
    #将测试图像转换为灰度图像，因为opencv人脸检测器需要灰度图像
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    #检测多尺度图像，返回值是一张脸部区域信息的列表（x,y,宽,高）
    faces = face_cascade.detectMultiScale(
        gray,
        scaleFactor=1.2,
        minNeighbors=5,
        minSize=(20, 20) #人脸的最小范围，如果比20*20像素小忽略
        )
 
    # 如果未检测到面部，则返回原始图像
    if (len(faces) == 0):
        return None, None 
    #目前假设只有一张脸，xy为左上角坐标，wh为矩形的宽高
    (x, y, w, h) = faces[0]
    #返回图像的正面部分
    return gray[y:y + w, x:x + h], faces[0]


def Renew_face_model():
    #调用prepare_training_data（）函数
    faces, labels = prepare_training_data("face_data")
     
    #创建LBPH识别器并开始训练，当然也可以选择Eigen或者Fisher识别器
    face_recognizer = cv2.face.LBPHFaceRecognizer_create()
    face_recognizer.train(faces, np.array(labels))

    #safe model
    face_recognizer.save("face_model.xml")
    return face_recognizer


def prepare_training_data(data_folder_path):
    # 获取数据文件夹中的目录（每个主题的一个目录）
    dirs = os.listdir(data_folder_path)
    #print(dirs)
    # 两个列表分别保存所有的脸部和标签
    faces = []
    labels = []
    # 浏览每个目录并访问其中的图像
    for dir_name in dirs:
        # dir_name(str类型)即标签
        label = int(dir_name)
        # 建立包含当前主题主题图像的目录路径
        subject_dir_path = data_folder_path + "/" + dir_name
        print(subject_dir_path+"\n")
        # 获取给定主题目录内的图像名称
        subject_images_names = os.listdir(subject_dir_path)
        # 浏览每张图片并检测脸部，然后将脸部信息添加到脸部列表faces[]
        for image_name in subject_images_names:
            # 建立图像路径
            image_path = subject_dir_path + "/" + image_name
            # 读取图像
            image = cv2.imread(image_path)
            # 显示图像0.1s
            cv2.imshow("Training on image...", image)
            cv2.waitKey(10)
            # 检测脸部
            face, rect, no_mask = detect_face(image)
            # 我们忽略未检测到的脸部
            if face is not None:
                #将脸添加到脸部列表并添加相应的标签
                faces.append(face)
                labels.append(label)
 
    cv2.waitKey(1)
    cv2.destroyAllWindows()
    #最终返回值为人脸和标签列表
    return faces, labels

#while adding people to the dataSet, capture ten pictures and
#then save into a folder
def Capture_Image():
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("cant initial camera")
    
    #make new folder to save face data
    path = os.getcwd()
    print(path)
    ls = os.listdir(path+"/face_data")
    #print(max(ls))
    new_folder = str(int(max(ls))+1)
    os.mkdir(path+"/face_data/"+new_folder)
    
    cnt = 0
    nimg = 0
    while True:
        ret, frame = cap.read()
        
        if not ret:
            print("cant obtain image")
            cap.release();
            break
        
        try:
            img = frame.copy()
            face, rect = detect_face(img)
            draw_rectangle(img, rect)
            if cnt == 10:
                cv2.imwrite("face_data/"+new_folder+"/"+str(nimg)+".jpg", frame)
                nimg = nimg + 1
                cnt = 0
            else:
                cnt = cnt + 1 
        except Exception as err:
            
            print(err)
        if nimg == 16:
            break
        draw_text(img, "Obtaining picture...", 20, 30, "GREEN")
        draw_text(img, "Please keep only one person ", 20, 60, "GREEN")
        draw_text(img, "in the camera", 20, 90, "GREEN")
        # cv2.imshow("picture", img)
        # cv2.waitKey(10)
        #cv2.imshow("display", frame)
    cap.release();
    cv2.destroyAllWindows()   
    
#change the existing people's image   
def Change_Image(filename):
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("cant initial camera")
    
    cnt = 0
    nimg = 0
    while True:
        ret, frame = cap.read()
        
        if not ret:
            print("cant obtain image")
            cap.release()
            break
        try:
            img = frame.copy()
            face, rect = detect_face(img)
            draw_rectangle(img, rect)
            if cnt == 10:
                cv2.imwrite("face_data/"+str(filename)+"/"+str(nimg)+".jpg", frame)
                nimg = nimg + 1
                cnt = 0
            else:
                cnt = cnt + 1 
        except Exception as err:
            print(err)
        if nimg == 16:
            break
        draw_text(img, "Obtaining picture...", 20, 30, "GREEN")
        draw_text(img, "Please keep only one person ", 20, 60, "GREEN")
        draw_text(img, "in the camera", 20, 90, "GREEN")

        cv2.imshow("picture", img)
        cv2.waitKey(10)
        
    cap.release()
    print('training done')
    cv2.destroyAllWindows()

def Renew_face_model():

    #调用prepare_training_data（）函数
    faces, labels = prepare_training_data("face_data")
     
    #创建LBPH识别器并开始训练，当然也可以选择Eigen或者Fisher识别器
    face_recognizer = cv2.face.LBPHFaceRecognizer_create()
    face_recognizer.train(faces, np.array(labels))

    #safe model
    face_recognizer.save("face_model.xml")
    return face_recognizer


def prepare_training_data(data_folder_path):
    # 获取数据文件夹中的目录（每个主题的一个目录）
    dirs = os.listdir(data_folder_path)
    #print(dirs)
    # 两个列表分别保存所有的脸部和标签
    faces = []
    labels = []
    # 浏览每个目录并访问其中的图像
    for dir_name in dirs:
        # dir_name(str类型)即标签
        label = int(dir_name)
        # 建立包含当前主题主题图像的目录路径
        subject_dir_path = data_folder_path + "/" + dir_name
        print(subject_dir_path+"\n")
        # 获取给定主题目录内的图像名称
        subject_images_names = os.listdir(subject_dir_path)
        # 浏览每张图片并检测脸部，然后将脸部信息添加到脸部列表faces[]
        for image_name in subject_images_names:
            # 建立图像路径
            image_path = subject_dir_path + "/" + image_name
            # 读取图像
            image = cv2.imread(image_path)
            # 显示图像0.1s
            # cv2.imshow("Training on image...", image)
            # cv2.waitKey(100)
            # 检测脸部
            face, rect = detect_face(image)
            # 我们忽略未检测到的脸部
            if face is not None:
                #将脸添加到脸部列表并添加相应的标签
                faces.append(face)
                labels.append(label)
 
    cv2.waitKey(1)
    cv2.destroyAllWindows()
    #最终返回值为人脸和标签列表
    return faces, labels


def renew_main():    
    Capture_Image()
    Renew_face_model()


def change_main(i):
    Change_Image(i)
    Renew_face_model()


def remove_face_data():
    #xxxxzzxxxxxxxxxxx
    dirs = os.listdir("face_data")
    for dir_name in dirs:
        # dir_name(str类型)即标签
        label = int(dir_name)
        if label > 3:
            imgs = os.listdir("face_data/"+str(label))
            for img in imgs:
                os.remove('face_data/'+str(label)+'/'+img)
            os.rmdir('face_data/'+str(label))
    Renew_face_model()

#if __name__ == '__main__':
#    renew_main()    
bt_add = tk.Button(window,text='录入人脸数据',height=2,width=18,command=renew_main)
bt_add.place(x=50,y=20)
bt_change_wht = tk.Button(window,text='更新吴宏涛数据',height=2,width=18,command=lambda: change_main(2))
bt_change_wht.place(x=50,y=70)
bt_change_zyn = tk.Button(window,text='更新郑轶楠数据',height=2,width=18,command=lambda: change_main(1))
bt_change_zyn.place(x=50,y=120)
bt_change_ozc = tk.Button(window,text='更新区梓川数据',height=2,width=18,command=lambda: change_main(3))
bt_change_ozc.place(x=50,y=170)
bt_rm = tk.Button(window,text='remove face data',height=2,width=18,command=remove_face_data)
bt_rm.place(x=50,y=220)

window.mainloop()
